#include <ros/ros.h>
#include <sensor_msgs/Image.h>
#include <sensor_msgs/CameraInfo.h>
#include <cv_bridge/cv_bridge.h>
#include <geometry_msgs/Pose.h>
#include <opencv2/opencv.hpp>
#include <opencv2/calib3d.hpp>
#include <tf2/LinearMath/Quaternion.h>
#include <tf2/LinearMath/Matrix3x3.h>

class ChessboardDetector {
public:
  ChessboardDetector() : nh_("~"), has_camera_info_(false) {
    // 从参数服务器获取配置
    nh_.param<int>("rows", chessboard_rows_, 12);
    nh_.param<int>("cols", chessboard_cols_, 9);
    nh_.param<double>("marker_size", square_size_, 0.03);
    nh_.param<bool>("image_is_rectified", image_is_rectified_, true);

    // 初始化3D物体点
    generateObjectPoints();

    // 初始化ROS订阅/发布
    ROS_INFO("initial");
    camera_info_sub_ = nh_.subscribe("/camera_info", 1, 
                                    &ChessboardDetector::cameraInfoCallback, this);
    image_sub_ = nh_.subscribe("/image", 1, 
                              &ChessboardDetector::imageCallback, this);
    pose_pub_ = nh_.advertise<geometry_msgs::Pose>("pose", 10);
  }

private:
  void generateObjectPoints() {
    obj_points_.clear();
    for(int i = 0; i < chessboard_cols_; ++i) {
      for(int j = 0; j < chessboard_rows_; ++j) {
        obj_points_.emplace_back(cv::Point3f(
          j * square_size_, i * square_size_, 0.0f));
      }
    }
    ROS_INFO_STREAM("Generated obj_points[0]: (" 
        << obj_points_[0].x << ", " << obj_points_[0].y << ")");
  }

  void cameraInfoCallback(const sensor_msgs::CameraInfo::ConstPtr& msg) {
    // if(has_camera_info_) return;

    // 解析相机内参矩阵
    camera_matrix_ = cv::Mat(3, 3, CV_64F, const_cast<double*>(msg->K.data())).clone();
    ROS_INFO("Camera parameters received 0");
    // 解析畸变系数
    if(image_is_rectified_) {
      dist_coeffs_ = cv::Mat::zeros(5, 1, CV_64F);
    } else {
      dist_coeffs_ = cv::Mat(msg->D).clone();
    }

    has_camera_info_ = true;
    ROS_INFO("Camera parameters received");
  }

  void imageCallback(const sensor_msgs::ImageConstPtr& msg) {
    if(!has_camera_info_) return;

    try {
      // 转换ROS图像到OpenCV格式
      cv_bridge::CvImagePtr cv_ptr = cv_bridge::toCvCopy(msg, "bgr8");
      cv::Mat gray;
      cv::cvtColor(cv_ptr->image, gray, cv::COLOR_BGR2GRAY);

      // 检测棋盘格角点
      std::vector<cv::Point2f> corners;
      bool found = cv::findChessboardCorners(
        gray, cv::Size(chessboard_rows_, chessboard_cols_), corners,
        cv::CALIB_CB_ADAPTIVE_THRESH | cv::CALIB_CB_NORMALIZE_IMAGE);

      if(found) {
        // 亚像素优化
        cv::cornerSubPix(gray, corners, cv::Size(11,11), cv::Size(-1,-1),
                        cv::TermCriteria(cv::TermCriteria::EPS + cv::TermCriteria::MAX_ITER, 30, 0.01));

        // 计算位姿
        cv::Mat rvec, tvec;
        cv::solvePnP(obj_points_, corners, camera_matrix_, dist_coeffs_, rvec, tvec);

        // 转换旋转向量到四元数
        cv::Mat R;
        cv::Rodrigues(rvec, R);
        tf2::Matrix3x3 tf_rot(
        R.at<double>(0,0), R.at<double>(0,1), R.at<double>(0,2),
        R.at<double>(1,0), R.at<double>(1,1), R.at<double>(1,2),
        R.at<double>(2,0), R.at<double>(2,1), R.at<double>(2,2)
        );
        tf2::Quaternion quat;
        tf_rot.getRotation(quat);

        // 发布位姿
        geometry_msgs::Pose pose_msg;
        pose_msg.position.x = tvec.at<double>(0);
        pose_msg.position.y = tvec.at<double>(1);
        pose_msg.position.z = tvec.at<double>(2);
        pose_msg.orientation.x = quat.x();
        pose_msg.orientation.y = quat.y();
        pose_msg.orientation.z = quat.z();
        pose_msg.orientation.w = quat.w();
        pose_pub_.publish(pose_msg);

        // 可视化（可选）
        cv::drawChessboardCorners(cv_ptr->image, 
                                cv::Size(chessboard_rows_, chessboard_cols_), 
                                corners, found);
        cv::imshow("Chessboard Detection", cv_ptr->image);
        cv::waitKey(1);
      }
    } catch (cv_bridge::Exception& e) {
      ROS_ERROR("cv_bridge exception: %s", e.what());
    }
  }

  ros::NodeHandle nh_;
  ros::Subscriber camera_info_sub_, image_sub_;
  ros::Publisher pose_pub_;

  int chessboard_rows_, chessboard_cols_;
  double square_size_;
  bool image_is_rectified_;
  bool has_camera_info_;

  cv::Mat camera_matrix_, dist_coeffs_;
  std::vector<cv::Point3f> obj_points_;
};

int main(int argc, char** argv) {
  ros::init(argc, argv, "chessboard_detector_cpp");
  ChessboardDetector detector;
  ros::spin();
  return 0;
}